# 代码
# WindowWidth = 640
# WindowHeight = 640
# SelectClass = '40070001'

# 以上三个参数是需要截图的窗口大小，和需要筛选迁移的标签
# 来源：https://blog.csdn.net/weixin_42112760/article/details/130537497


"""
选择裁剪图片窗口的大小，并且筛选遍历各个窗口内对应筛选种类的标签，对标签进行VOC格式存储
"""

import xml.etree.ElementTree as ET
import os
from PIL import Image


def check_path(Path):
    """
    # 修改格式和确认路径存在
    """
    if not os.path.exists(Path):
        os.makedirs(Path)
        print('Path does not exist. Build new Path:{}'.format(Path))
    if Path[-1] == '/':
        Path += '/'
    return


def pretty_xml(element, indent, newline, level=0):  # elemnt为传进来的Elment类，参数indent用于缩进，newline用于换行
    if element:  # 判断element是否有子元素
        if (element.text is None) or element.text.isspace():  # 如果element的text没有内容
            element.text = newline + indent * (level + 1)
        else:
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
            # else:  # 此处两行如果把注释去掉，Element的text也会另起一行
            # element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
    temp = list(element)  # 将element转成list
    for subelement in temp:
        if temp.index(subelement) < (len(temp) - 1):  # 如果不是list的最后一个元素，说明下一个行是同级别元素的起始，缩进应一致
            subelement.tail = newline + indent * (level + 1)
        else:  # 如果是list的最后一个元素， 说明下一行是母元素的结束，缩进应该少一个
            subelement.tail = newline + indent * level
        pretty_xml(subelement, indent, newline, level=level + 1)  # 对子元素进行递归操作


def new_voc_xml(folder_name, filename, path_name, size_list):
    """
    创建VOC_xml base框架
    folder_name: data_91_101
    xml_name: 091_0053.jpg
    path_name: /Users/miag/Documents/项目/code/data_91_101/091_0053.jpg
    size_list: [width, height, depth]
    """
    new_tree = ET.Element("annotation")
    build_new_node(new_tree, "folder", folder_name)
    build_new_node(new_tree, "filename", filename)
    build_new_node(new_tree, "path", path_name)
    size = ET.SubElement(new_tree, "size")
    build_new_node(size, "width", str(size_list[0]))
    build_new_node(size, "height", str(size_list[1]))
    build_new_node(size, "depth", str(size_list[2]))
    build_new_node(new_tree, "segmented", "0")
    return new_tree


def build_new_node(root, tag, text):
    child_node = ET.SubElement(root, tag)
    child_node.text = text


def is_update_tree(new_tree, name, xmin, ymin, xmax, ymax, a0, b0, a1, b1):
    """
    对选中区域内对应对筛选目标进行坐标转换，并更新在新的xml树上
    """
    if xmin >= a0 and ymin >= b0:
        update_xmin = xmin - a0
        update_ymin = ymin - b0
        if xmax <= a1 and ymax <= b1:
            update_xmax = xmax - a0
            update_ymax = ymax - b0
        elif abs(a1 - xmin) <= 32 and abs(b1 - ymin) <= 32:
            update_xmax = a1
            update_ymax = b1
        else:
            return new_tree
        object = ET.SubElement(new_tree, "object")
        build_new_node(object, "name", name)
        bndbox = ET.SubElement(object, "bndbox")
        build_new_node(bndbox, "xmin", str(update_xmin))
        build_new_node(bndbox, "ymin", str(update_ymin))
        build_new_node(bndbox, "xmax", str(update_xmax))
        build_new_node(bndbox, "ymax", str(update_ymax))
        new_tree.append(object)
        # print("find" + name)
    return new_tree


def get_bndbox(object):
    bndbox = object.find('bndbox')  # 访问bndbox节点
    xmin = bndbox.find("xmin").text  # 从bndbox节点中找到xmin的值
    xmax = bndbox.find("xmax").text
    ymin = bndbox.find("ymin").text  # 从bndbox节点中找到xmin的值
    ymax = bndbox.find("ymax").text
    # print(type(ymax))
    return int(xmin), int(ymin), int(xmax), int(ymax)


ImgPath = r'imgs/'
AnnoPath = r'xmls/'
ProcessedImgPath = r'slipOutput/JPEGImages/'
ProcessedAnnoPath = r'slipOutput/Annotations/'

# 修改格式和确认路径存在
check_path(ImgPath)
check_path(AnnoPath)
check_path(ProcessedImgPath)
check_path(ProcessedAnnoPath)

ImageList = os.listdir(ImgPath)
ImageList = sorted(ImageList)
ImageType = ['jpg', 'png', 'jpeg']

WindowWidth = 640
WindowHeight = 640
SelectClass = '40070001'

for image in ImageList:
    # ImageTitle, ext = os.path.splitext(image)
    ImageTitle, ext = image.split('.')
    if not (ext in ImageType):
        ImageList.remove(image)
        continue
    ImageFile = ImgPath + image
    XmlFile = AnnoPath + ImageTitle + '.xml'

    # 打开对应的xml文件
    tree = ET.parse(XmlFile)
    root = tree.getroot()  # 使用getroot()获取根节点，得到的是一个Element对象

    # 读取输入图片
    ReadImage = Image.open(ImageFile)
    width, height = ReadImage.size
    depth = len(ReadImage.mode)
    print("Input image {}: width{}, height{}, depth{}".format(image, width, height, depth))

    # 标记坐标
    position_w, position_h = 0, 0
    # 循环每个区域 a(0,8000),b(0,6000)， 起点在距离边界0处
    for a0 in range(0, (width - WindowWidth), WindowWidth):
        for b0 in range(0, (height - WindowHeight), WindowHeight):
            a1 = a0 + WindowWidth - 1
            b1 = b0 + WindowHeight - 1

            Annotations = ProcessedAnnoPath.split('/')[-2]  # Annotations
            filename = ImageTitle + '_{}{}.{}'.format(position_w, position_h, ext)
            path_name = ProcessedImgPath + filename
            size_list = [WindowWidth, WindowHeight, depth]
            new_tree = new_voc_xml(Annotations, filename, path_name, size_list)

            for object in root.findall('object'):
                name = object.find('name').text  # 访问Element文本
                if name == SelectClass or SelectClass is None:
                    # print('------> get ' + name )
                    xmin, ymin, xmax, ymax = get_bndbox(object)
                    new_tree = is_update_tree(new_tree, name, xmin, ymin, xmax, ymax, a0, b0, a1, b1)

            new_xml_tree = ET.ElementTree(new_tree)
            new_root = new_xml_tree.getroot()  # 得到根元素，Element类
            if new_root.findall('object'):
                print('object counts: {}'.format(len(new_root.findall('object'))))
                pretty_xml(new_root, '\t', '\n')  # 执行美化方法
                xml_file = ProcessedAnnoPath + ImageTitle + '_{}{}.xml'.format(position_w, position_h)
                new_xml_tree.write(xml_file)
                cut_area = ReadImage.crop([a0, b0, a1, b1])
                cut_area.save(path_name)
                print('Succeed to save xml and img: ' + filename)

            new_tree = None
            position_h += 1
        position_w += 1
        position_h =0
    # break